from modelscope import AutoModel, AutoModelForCausalLM, AutoTokenizer, AutoImageProcessor
from modelscope import Qwen2VLForConditionalGeneration, AutoProcessor
import torch

def load_model(models_dir, model_name):
    processor = None
    if model_name == "MiniCPM-Llama3-V-2_6-int4":
        tokenizer = AutoTokenizer.from_pretrained(models_dir["MiniCPM-Llama3-V-2_6-int4"], trust_remote_code=True)
        model = AutoModel.from_pretrained(models_dir["MiniCPM-Llama3-V-2_6-int4"], trust_remote_code=True)
    elif model_name == "deepseek-vl2-tiny":
        tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/deepseek-vl2-tiny", trust_remote_code=True)
        model = AutoModelForCausalLM.from_pretrained("deepseek-ai/deepseek-vl2-tiny", trust_remote_code=True)
    elif model_name == "InternVL2_5-2B":
        tokenizer = AutoTokenizer.from_pretrained(models_dir["InternVL2_5-2B"], trust_remote_code=True)
        model = AutoModel.from_pretrained(
            models_dir["InternVL2_5-2B"],
            torch_dtype=torch.bfloat16,
            low_cpu_mem_usage=True,
            use_flash_attn=True,
            trust_remote_code=True,
            device_map="auto")
    elif model_name == "Qwen2_VL-2B":
        model = Qwen2VLForConditionalGeneration.from_pretrained(
            models_dir["Qwen2_VL-2B"], torch_dtype="auto", device_map="auto"
        )
        processor = AutoProcessor.from_pretrained(models_dir["Qwen2_VL-2B"])
        tokenizer = None
        
    elif model_name == "GLM-Edge-2B":
        tokenizer = AutoTokenizer.from_pretrained(models_dir["GLM-Edge-2B"], trust_remote_code=True)
        model = AutoModelForCausalLM.from_pretrained(models_dir["GLM-Edge-2B"], 
                                                     torch_dtype=torch.bfloat16, 
                                                     trust_remote_code=True, 
                                                     device_map="auto")
        processor = AutoImageProcessor.from_pretrained(models_dir["GLM-Edge-2B"], trust_remote_code=True, device_map="auto")

    model.eval()

    return model, tokenizer, processor